#!/usr/bin/env python3
import rospy
import numpy as np
from sensor_msgs.msg import Image, CameraInfo
from std_msgs.msg import Float32MultiArray
from geometry_msgs.msg import PointStamped
from cv_bridge import CvBridge

class Object3DDetector:
    def __init__(self):
        rospy.init_node('object_3d_detector', anonymous=True)

        self.bridge = CvBridge()

        # 相机内参
        self.color_K = None
        self.depth_K = None

        # 最新深度图
        self.latest_depth = None

        # 滑动平均历史
        self.history = {}  # key: cls_id, value: list of (X,Y,Z)
        self.window_size = 5  # 滑动平均帧数

        # 订阅相机信息
        rospy.Subscriber("/camera/color/camera_info", CameraInfo, self.color_info_callback)
        rospy.Subscriber("/camera/depth/camera_info", CameraInfo, self.depth_info_callback)

        # 订阅深度图
        rospy.Subscriber("/camera/depth/image_raw", Image, self.depth_callback)

        # 订阅 YOLO 检测结果
        rospy.Subscriber("/yolo/detections", Float32MultiArray, self.yolo_callback)

        # 发布三维坐标
        self.point_pub = rospy.Publisher("/yolo/object_3d", PointStamped, queue_size=1)

    # ---------------- 回调 ----------------
    def color_info_callback(self, msg):
        self.color_K = np.array(msg.K).reshape(3,3)

    def depth_info_callback(self, msg):
        self.depth_K = np.array(msg.K).reshape(3,3)

    def depth_callback(self, msg):
        self.latest_depth = self.bridge.imgmsg_to_cv2(msg, "passthrough")  # mm

    def yolo_callback(self, msg):
        if self.latest_depth is None or self.color_K is None:
            return

        data = msg.data
        # 每个目标 6 个元素: cls, conf, cx, cy, w, h
        for i in range(0, len(data), 6):
            cls_id = int(data[i])
            conf = data[i+1]
            cx = int(data[i+2])
            cy = int(data[i+3])
            w = data[i+4]
            h = data[i+5]

            # 防止越界
            h_max, w_max = self.latest_depth.shape
            cx = np.clip(cx, 1, w_max-2)
            cy = np.clip(cy, 1, h_max-2)

            # 深度中值滤波: 取目标中心 3x3 区域
            roi = self.latest_depth[cy-1:cy+2, cx-1:cx+2]
            z_mm = np.median(roi)
            if z_mm == 0:
                continue  # 忽略无效深度
            z = z_mm / 1000.0  # 转米

            fx = self.color_K[0,0]
            fy = self.color_K[1,1]
            cx_k = self.color_K[0,2]
            cy_k = self.color_K[1,2]

            # 像素 -> 相机坐标
            X = (cx - cx_k) * z / fx
            Y = (cy - cy_k) * z / fy
            Z = z

            # 滑动平均滤波
            key = cls_id  # 可以根据需要加 bbox 索引区分不同目标
            if key not in self.history:
                self.history[key] = []
            self.history[key].append((X,Y,Z))
            if len(self.history[key]) > self.window_size:
                self.history[key].pop(0)

            X_avg = np.mean([p[0] for p in self.history[key]])
            Y_avg = np.mean([p[1] for p in self.history[key]])
            Z_avg = np.mean([p[2] for p in self.history[key]])

            # 发布三维坐标
            point_msg = PointStamped()
            point_msg.header.stamp = rospy.Time.now()
            point_msg.header.frame_id = "camera_color_optical_frame"
            point_msg.point.x = X_avg
            point_msg.point.y = Y_avg
            point_msg.point.z = Z_avg
            self.point_pub.publish(point_msg)

    # ---------------- 运行 ----------------
    def run(self):
        rospy.spin()

if __name__ == "__main__":
    node = Object3DDetector()
    node.run()
